

import * as tf from '@tensorflow/tfjs';
import { Conv2D } from '@tensorflow/tfjs-layers/dist/layers/convolutional';
import { DepthwiseConv2D } from '@tensorflow/tfjs-layers/dist/layers/convolutional_depthwise';
import { BatchNormalization } from '@tensorflow/tfjs-layers/dist/layers/normalization';



import { GlobalAveragePooling2D, AveragePooling2D } from '@tensorflow/tfjs-layers/dist/layers/pooling';
import { Dense, Activation } from '@tensorflow/tfjs-layers/dist/layers/core';
import { Kwargs } from '@tensorflow/tfjs-layers/dist/types';
import { ContainerArgs } from '@tensorflow/tfjs-layers/dist/engine/container';

const IMAGE_SIZE = 224;

function h_sigmoid(x: tf.Tensor) {
    return tf.relu6(x.add(3)).div(6);
}


function h_swish(x: tf.Tensor) {
    return h_sigmoid(x).mul(x);
}

class SEBlock extends tf.layers.Layer {
    pool: GlobalAveragePooling2D;
    fc1: Dense;
    fc2: Dense;
    getClassName(): string {
        return "SEBlock";
    }
    constructor(input_channels: number, r = 16) {
        super({});
        this.pool = tf.layers.globalAveragePooling2d({});
        this.fc1 = tf.layers.dense({ units: Math.floor(input_channels / r) })
        this.fc2 = tf.layers.dense({ units: input_channels });
    }


    call(inputs: tf.Tensor, kwargs: Kwargs): tf.Tensor | tf.Tensor[] {
        let branch = this.pool.call(inputs, kwargs);
        branch = this.fc1.call(branch, kwargs);
        if (branch instanceof tf.Tensor) {
            branch = tf.relu(branch)
        }
        branch = this.fc2.call(branch, kwargs)

        if (branch instanceof tf.Tensor) {
            branch = h_sigmoid(branch)
            branch = tf.expandDims(branch, 1)
            branch = tf.expandDims(branch, 1)
            let output = inputs.mul(branch)
            return output
        }
        return [];
    }
}



class BottleNeck extends tf.layers.Layer {
    in_size: number;
    out_size: number;
    stride: number | number[];
    is_se_existing: boolean;
    NL: string;

    conv1: Conv2D;
    conv2: Conv2D;
    bn1: BatchNormalization;
    bn2: BatchNormalization;
    bn3: BatchNormalization;
    dwconv: DepthwiseConv2D;
    se: SEBlock;
    linear: Activation;
    getClassName() {
        return "BottleNeck";
    }

    constructor(in_size: number, exp_size: number, out_size: number, stride: number | number[], is_se_existing: boolean, NL: string, k: number) {
        super({});

        this.stride = stride;
        this.in_size = in_size;
        this.out_size = out_size;
        this.is_se_existing = is_se_existing;
        this.NL = NL;
        this.conv1 = tf.layers.conv2d({
            filters: exp_size,
            kernelSize: [1, 1],
            strides: 1,
            padding: "same"
        });
        this.bn1 = tf.layers.batchNormalization();
        this.dwconv = tf.layers.depthwiseConv2d({
            kernelSize: [k, k],
            strides: stride,
            padding: "same"
        });

        this.bn2 = tf.layers.batchNormalization();
        this.se = new SEBlock(exp_size);
        this.conv2 = tf.layers.conv2d({
            filters: out_size,
            kernelSize: [1, 1],
            strides: 1,
            padding: "same"
        });
        this.bn3 = tf.layers.batchNormalization();

        this.linear = tf.layers.activation({ activation: 'linear' });
    }

    call(inputs: tf.Tensor, /* training = 'None',*/ kwargs: Kwargs) {
        let x = this.conv1.call(inputs, kwargs);
        x = this.bn1.call(x, kwargs);

        if (this.NL == "HS" && x instanceof tf.Tensor) {
            x = h_swish(x);
        } else if (this.NL == "RE" && x instanceof tf.Tensor) {
            x = tf.relu6(x);
        }
        x = this.dwconv.call(x, kwargs)
        x = this.bn2.call(x, kwargs)

        if (this.NL == "HS" && x instanceof tf.Tensor) {
            x = h_swish(x)

        } else if (this.NL == "RE" && x instanceof tf.Tensor) {
            x = tf.relu6(x)
        }


        if (this.is_se_existing && x instanceof tf.Tensor) {
            x = this.se.call(x, kwargs)
        }
        x = this.conv2.call(x, kwargs)
        x = this.bn3.call(x, kwargs)
        x = this.linear.call(x, kwargs)

        if (this.stride == 1 && this.in_size == this.out_size) {
            //x = tf.layers.add([x, inputs])
            x = x.concat(inputs);
        }

        return x
    }
}

const NUM_CLASSES = 10;

export class MobileNetV3 {
    model: tf.Sequential;
    constructor() {
        let model = tf.sequential();
        this.model = model;
        model.add(tf.layers.conv2d({
            filters: 16,
            kernelSize: [3, 3],
            strides: 2,
            padding: "same",
            inputShape: [224, 224, 3]
        }));

        model.add(tf.layers.batchNormalization());


        model.add(new BottleNeck(16, 16, 16, 2, true, "RE", 3))
        model.add(new BottleNeck(16, 72, 24, 2, false, "RE", 3))
        model.add(new BottleNeck(24, 88, 24, 1, false, "RE", 3))
        model.add(new BottleNeck(24, 96, 40, 2, true, "HS", 5))
        model.add(new BottleNeck(40, 240, 40, 1, true, "HS", 5))
        model.add(new BottleNeck(40, 240, 40, 1, true, "HS", 5))
        model.add(new BottleNeck(40, 120, 48, 1, true, "HS", 5))
        model.add(new BottleNeck(48, 144, 48, 1, true, "HS", 5))
        model.add(new BottleNeck(48, 288, 96, 2, true, "HS", 5))
        model.add(new BottleNeck(96, 576, 96, 1, true, "HS", 5))
        model.add(new BottleNeck(96, 576, 96, 1, true, "HS", 5))

        model.add(tf.layers.conv2d({
            filters: 576,
            kernelSize: [1, 1],
            strides: 1,
            padding: "same"
        }));

        model.add(tf.layers.batchNormalization());


        model.add(tf.layers.averagePooling2d({ poolSize: [7, 7], strides: 1 }));
        model.add(tf.layers.conv2d({
            filters: 1280,
            kernelSize: [1, 1],
            strides: 1,
            padding: "same"
        }));
        model.add(tf.layers.conv2d({
            filters: NUM_CLASSES,
            kernelSize: [1, 1],
            strides: 1,
            padding: "same",
            activation: 'softmax'
        }));
    }

    infer(
        img: tf.Tensor | ImageData | HTMLImageElement | HTMLCanvasElement |
            HTMLVideoElement,
        embedding = false): tf.Tensor {
        return tf.tidy(() => {
            if (!(img instanceof tf.Tensor)) {
                img = tf.browser.fromPixels(img);
            }

            // Normalize the image from [0, 255] to [inputMin, inputMax].
            const normalized: tf.Tensor3D = tf.add(
                tf.mul(tf.cast(img, 'float32'), 1.0 / 255), 0);

            // Resize the image to
            let resized = normalized;
            if (img.shape[0] !== IMAGE_SIZE || img.shape[1] !== IMAGE_SIZE) {
                const alignCorners = true;
                resized = tf.image.resizeBilinear(
                    normalized, [IMAGE_SIZE, IMAGE_SIZE], alignCorners);
            }

            // Reshape so we can pass it to predict.
            const batched = tf.reshape(resized, [-1, IMAGE_SIZE, IMAGE_SIZE, 3]);

            let result: tf.Tensor2D;

            if (embedding) {
                //const embeddingName = EMBEDDING_NODES[this.version];
                const embeddingName = "1";
                const a = this.model.fit(batched, tf.tensor1d([1, 0, 0]));
                const internal = this.model.execute(batched, embeddingName) as tf.Tensor4D;
                result = tf.squeeze(internal);
                //        result = tf.squeeze(internal, [1, 2]);
            } else {
                const logits1001 = this.model.predict(batched) as tf.Tensor2D;
                // Remove the very first logit (background noise).
                result = tf.slice(logits1001, [0, 1], [-1, 1000]);
            }

            return result;
        });
    }

    async fit(
        img: tf.Tensor | ImageData | HTMLImageElement | HTMLCanvasElement |
            HTMLVideoElement, id: number) {

        if (!(img instanceof tf.Tensor)) {
            img = tf.browser.fromPixels(img);
        }

        // Normalize the image from [0, 255] to [inputMin, inputMax].
        const normalized: tf.Tensor3D = tf.add(
            tf.mul(tf.cast(img, 'float32'), 1.0 / 255), 0);

        // Resize the image to
        let resized = normalized;
        if (img.shape[0] !== IMAGE_SIZE || img.shape[1] !== IMAGE_SIZE) {
            const alignCorners = true;
            resized = tf.image.resizeBilinear(
                normalized, [IMAGE_SIZE, IMAGE_SIZE], alignCorners);
        }

        // Reshape so we can pass it to predict.
        const batched = tf.reshape(resized, [-1, IMAGE_SIZE, IMAGE_SIZE, 3]);
        let zeros = tf.fill([1000], 0);
        let y = await zeros.buffer();
        y.set(id, 1);
        this.model.fit(batched, y.toTensor());
    }

}


export class MobileNetV3Small extends tf.LayersModel {
    conv1: Conv2D;
    conv2: Conv2D;
    conv3: Conv2D;
    conv4: Conv2D;
    bn1: BatchNormalization;
    bn2: BatchNormalization;

    bneck1: BottleNeck;
    bneck2: BottleNeck;
    bneck3: BottleNeck;
    bneck4: BottleNeck;
    bneck5: BottleNeck;
    bneck6: BottleNeck;
    bneck7: BottleNeck;
    bneck8: BottleNeck;
    bneck9: BottleNeck;
    bneck10: BottleNeck;
    bneck11: BottleNeck;

    avgpool: AveragePooling2D;

    constructor(args: ContainerArgs) {
        super(args);
        this.conv1 = tf.layers.conv2d({
            filters: 16,
            kernelSize: [3, 3],
            strides: 2,
            padding: "same",
            inputShape: [224, 224, 3]
        });

        this.bn1 = tf.layers.batchNormalization();

        this.bneck1 = new BottleNeck(16, 16, 16, 2, true, "RE", 3)
        this.bneck2 = new BottleNeck(16, 72, 24, 2, false, "RE", 3)
        this.bneck3 = new BottleNeck(24, 88, 24, 1, false, "RE", 3)
        this.bneck4 = new BottleNeck(24, 96, 40, 2, true, "HS", 5)
        this.bneck5 = new BottleNeck(40, 240, 40, 1, true, "HS", 5)
        this.bneck6 = new BottleNeck(40, 240, 40, 1, true, "HS", 5)
        this.bneck7 = new BottleNeck(40, 120, 48, 1, true, "HS", 5)
        this.bneck8 = new BottleNeck(48, 144, 48, 1, true, "HS", 5)
        this.bneck9 = new BottleNeck(48, 288, 96, 2, true, "HS", 5)
        this.bneck10 = new BottleNeck(96, 576, 96, 1, true, "HS", 5)
        this.bneck11 = new BottleNeck(96, 576, 96, 1, true, "HS", 5)

        this.conv2 = tf.layers.conv2d({
            filters: 576,
            kernelSize: [1, 1],
            strides: 1,
            padding: "same"
        });

        this.bn2 = tf.layers.batchNormalization();


        this.avgpool = tf.layers.averagePooling2d({ poolSize: [7, 7], strides: 1 });
        this.conv3 = tf.layers.conv2d({
            filters: 1280,
            kernelSize: [1, 1],
            strides: 1,
            padding: "same"
        });
        this.conv4 = tf.layers.conv2d({
            filters: NUM_CLASSES,
            kernelSize: [1, 1],
            strides: 1,
            padding: "same",
            activation: 'softmax'
        });

    }

    call(inputs: tf.Tensor, kwargs: Kwargs) {
        let x = this.conv1.apply(inputs)
        x = this.bn1.apply(x, kwargs)
        x = h_swish(<tf.Tensor>x)
        x = this.bneck1.apply(x, kwargs)
        x = this.bneck2.apply(x, kwargs)
        x = this.bneck3.apply(x, kwargs)
        x = this.bneck4.apply(x, kwargs)
        x = this.bneck5.apply(x, kwargs)
        x = this.bneck6.apply(x, kwargs)
        x = this.bneck7.apply(x, kwargs)
        x = this.bneck8.apply(x, kwargs)
        x = this.bneck9.apply(x, kwargs)
        x = this.bneck10.apply(x, kwargs)
        x = this.bneck11.apply(x, kwargs)

        x = this.conv2.apply(x)
        x = this.bn2.apply(x, kwargs)
        x = h_swish(<tf.Tensor>x)
        x = this.avgpool.apply(x)
        x = this.conv3.apply(x)
        x = h_swish(<tf.Tensor>x)
        x = this.conv4.apply(x)

        return <tf.Tensor>x;
    }
}